/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
 *
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

 #include <vector>
 #include <cstdint>
 #include <cmath>
 #include <random>
 #include <iostream>
 #include <algorithm>
 #include "acl/acl.h"
 #include "elewise/elewise_host.h"
 #include "elewise/elewise_device.h" 
 #include "example_common.h"
 
 namespace {
 // 描述算子的输入输出以及临时计算资源
 using LogSigmoidOpTraits = ATVC::OpTraits<ATVC::OpInputs<float>, ATVC::OpOutputs<float>, ATVC::OpTemps<float, float,float>>;
 
 // 传入编译态参数ATVC::OpTraits
 template<class Traits>
 // 开发自定义函数名/类名
 struct LogSigmoidComputeFunc {
     // DataType模板参数，根据实际数据类型个数填写
     template<typename T, typename U>
     // 重载operator公有接口，提供给`ATVC::Kernel::EleWiseOpTemplate`调用
     __aicore__ inline void operator()(AscendC::LocalTensor<T> x,
                                       AscendC::LocalTensor<T> y,
                                       AscendC::LocalTensor<U> tempNegX,
                                      AscendC::LocalTensor<U> tempExp,
                                      AscendC::LocalTensor<U> tempLog)
     {
         // 开发调用AscendC Api自行实现计算仿函数
         uint32_t tiledCnt = y.GetSize(); // 进行单次基块计算的元素个数
        // tempNegX = -x
        AscendC::Muls(tempNegX, x, static_cast<T>(-1), tiledCnt);
        // tempExp = exp(-x)
        AscendC::Exp(tempExp, tempNegX, tiledCnt);
        // tempExp = 1 + exp(-x)
        AscendC::Adds(tempExp, tempExp, static_cast<T>(1.0), tiledCnt);
        // tempLog = log(1 + exp(-x))
        AscendC::Log(tempLog, tempExp, tiledCnt);
        // y = -log(1 + exp(-x))
        AscendC::Muls(y, tempLog, static_cast<T>(-1.0), tiledCnt);
     }
 };
 
 void InitializeData(int32_t eleNum, std::vector<float> &inputX, std::vector<float> &golden)
 {
     std::random_device rd;
     std::mt19937 gen(rd());
     std::uniform_real_distribution<float> dis(-10.0f, 10.0f);
 
     for (int i = 0; i < eleNum; ++i) {
         inputX[i] = dis(gen);
         golden[i] =-std::log1p(std::exp(-inputX[i])); 
     }
 }
 }
 

 template<class OpTraits>
 __global__ __aicore__ void LogSigmoidCustom(GM_ADDR x, GM_ADDR y, ATVC::EleWiseParam param)
 {
     KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY); // 控制算子执行时只启动Vector核
     auto op = ATVC::Kernel::EleWiseOpTemplate<LogSigmoidComputeFunc<OpTraits>>();
     op.Run(x, y, &param);          // 按照输入、输出、param的顺序传入Run函数中；OpTraits内部的ATVC::OpTemps将由EleWiseOpTemplate内部申请资源，开发无需关注
 }
 
 int main()
 {
     // init data
     int32_t eleNum = 8 * 2048;
     size_t inputByteSize = static_cast<size_t>(eleNum) * sizeof(float);
     size_t outputByteSize = static_cast<size_t>(eleNum) * sizeof(float);
 
     std::vector<float> inputX(eleNum);
     std::vector<float> golden(eleNum);
     InitializeData(eleNum, inputX, golden);
 
     ATVC::EleWiseParam param;
 
     // 计算输入为8*2048个float元素的LogSigmoid算子的运行态参数param
     if (!ATVC::Host::CalcEleWiseTiling<LogSigmoidOpTraits>(eleNum, param)) {
         printf("Elewise tiling error.");
         return -1;
     };
     // 初始化Acl资源与数据
     aclrtContext context;
     aclrtStream stream = nullptr;
     int32_t deviceId = 0;
     InitializeACL(context, stream, deviceId);
 
     uint8_t *yHost;
     uint8_t *xDevice;
     uint8_t *yDevice;
     CHECK_ACL(aclrtMalloc((void **)&xDevice, inputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
     CHECK_ACL(aclrtMemcpy(xDevice, inputByteSize, inputX.data(), inputByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
 
     CHECK_ACL(aclrtMallocHost((void **)(&yHost), outputByteSize));
     CHECK_ACL(aclrtMalloc((void **)&yDevice, outputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
 
     // 调用自定义的Kernel API, <<<>>>的BlockNum参数可通过param的TilingData获取
     LogSigmoidCustom<LogSigmoidOpTraits><<<param.tilingData.blockNum, nullptr, stream>>>(xDevice, yDevice, param);
 
     CHECK_ACL(aclrtSynchronizeStream(stream));
     CHECK_ACL(aclrtMemcpy(yHost, outputByteSize, yDevice, outputByteSize, ACL_MEMCPY_DEVICE_TO_HOST));
     std::vector<float> outputY(reinterpret_cast<float*>(yHost), reinterpret_cast<float*>(yHost) + eleNum);
 
     // 释放资源
     CHECK_ACL(aclrtFree(xDevice));
     CHECK_ACL(aclrtFree(yDevice));
     CHECK_ACL(aclrtFreeHost(yHost));
 
     CleanACL(stream, context, deviceId);
 
     if (!VerifyResults(golden, outputY)) {
         return -1;
     }
     printf("Accuracy verification passed.\n");
     return 0;
 }